'''
Created on 13.08.2012

@author: myaut
'''

import matplotlib.pyplot as plt
import matplotlib.figure as figure

from lrex.spec import Probe

annotate = False

class RQVizualizer:
    colors = ('#CC9933', '#FFFF66', '#339966', '#CCFFCC', 
                     '#9999FF', '#33FFCC')
    
    def __init__(self, trace):
        self.trace = trace
    
    def calc_width(self, rrq):        
        start_time = rrq.get_start_time()
        last_end_time = rrq.get_end_time()
        
        lvl = 1
        max_lvl = 1
        
        min_rq_time = last_end_time - start_time
        
        rq_list = [(lvl, rrq)]
        
        while rq_list:
            lvl, rq = rq_list.pop()
            
            end_time = rq.get_end_time()
            rq_time = end_time - rq.get_start_time()
            
            min_rq_time = min(min_rq_time, rq_time)
            last_end_time = max(last_end_time, end_time)
            max_lvl = max(max_lvl, lvl)
            
            rq_list.extend([(lvl + 1, srq) for srq in rq.sub_rq])
        
        total_time = last_end_time - start_time
        
        width = 0.4 * (total_time / float(min_rq_time))
        height = 1.6 * max_lvl
        
        return min(250, width), min(100, height) 
        
    
    def draw(self, rrq, fig, lvl = 0):
        last_rq_end = 0
        h = 0.4
        
        for pair in rrq.pairs:
            x = pair.get_start_time()
            y = lvl * h
            
            col = RQVizualizer.colors[lvl % len(RQVizualizer.colors)]
            
            if annotate:
                plt.annotate(pair.get_name(), xy=(x, y + h), rotation='vertical',
                             arrowprops=dict(arrowstyle="->", connectionstyle="arc"))
                
            plt.barh(y, pair.get_duration(), h, x, color = col)
        
        for rq in rrq.sub_rq:            
            # if rq.get_start_time() < last_rq_end:
            #    lvl += 1
            
            self.draw(rq, fig, lvl + 1)
            
            last_rq_end = rq.get_end_time()
    
    def vizualize(self, viz = lambda rq: True):
        for rq in reversed(self.trace.requests):
            if not viz(rq):
                continue
            
            fn = '/tmp/lrex/rq-' + '%04d' % rq.id + '.png'
            
            start = rq.get_start_time()
            end = rq.get_end_time()
            
            figsize = self.calc_width(rq)
            
            fig = plt.figure(figsize = figsize)
            
            self.draw(rq, fig)
            
            fig.savefig(fn)
    
    def vizualize_all(self, viz = lambda rq: True):
        fn = '/tmp/lrex/rq-all.png'
        
        # figsize = self.calc_width(self.trace.requests[0])
        fig = plt.figure(figsize = (20, 6))
        
        for rq in reversed(self.trace.requests):
            if not viz(rq):
                continue
            
            self.draw(rq, fig)
            
        fig.savefig(fn)